#!/bin/bash

echo "🔄 USB设备重新识别工具"
echo "========================"

# 函数：显示可用的USB主机控制器
show_usb_controllers() {
    echo "🔌 可用的USB主机控制器："
    echo "========================"
    lsusb -t 2>/dev/null | grep -E "Host|Root Hub" || {
        echo "使用lsusb查看USB设备树："
        lsusb
    }
    echo "========================"
}

# 函数：重新扫描SCSI总线（适用于USB存储设备）
rescan_scsi_bus() {
    echo "🔄 重新扫描SCSI总线..."
    
    # 设置超时时间（3分钟）
    local timeout_duration=180
    local start_time=$(date +%s)
    
    # 方法1: 使用rescan-scsi-bus.sh（如果可用）
    if command -v rescan-scsi-bus.sh >/dev/null 2>&1; then
        echo "使用rescan-scsi-bus.sh工具..."
        
        # 在后台执行扫描命令
        sudo rescan-scsi-bus.sh &
        local scan_pid=$!
        
        # 等待进程完成或超时
        while ps -p $scan_pid > /dev/null 2>&1; do
            local current_time=$(date +%s)
            local elapsed_time=$((current_time - start_time))
            
            if [ $elapsed_time -ge $timeout_duration ]; then
                echo "⏰ 扫描超时（超过3分钟），终止扫描进程..."
                sudo kill $scan_pid 2>/dev/null
                wait $scan_pid 2>/dev/null
                echo "⚠️  SCSI总线扫描超时，继续执行后续步骤..."
                return 1
            fi
            
            sleep 2
        done
        
        # 检查进程退出状态
        wait $scan_pid
        local exit_status=$?
        
        if [ $exit_status -eq 0 ]; then
            echo "✅ rescan-scsi-bus.sh 扫描完成"
            return 0
        else
            echo "❌ rescan-scsi-bus.sh 扫描失败"
            return $exit_status
        fi
    fi
    
    # 方法2: 手动扫描所有SCSI主机
    echo "手动扫描SCSI主机..."
    local found=0
    
    # 记录开始扫描的时间
    start_time=$(date +%s)
    
    for host in /sys/class/scsi_host/host*/scan; do
        # 检查是否超时
        local current_time=$(date +%s)
        local elapsed_time=$((current_time - start_time))
        
        if [ $elapsed_time -ge $timeout_duration ]; then
            echo "⏰ 扫描超时（超过3分钟），停止扫描..."
            echo "⚠️  SCSI总线扫描超时，继续执行后续步骤..."
            return 1
        fi
        
        if [ -w "$host" ]; then
            echo "- - -" | sudo tee "$host" >/dev/null 2>&1
            echo "已扫描: $(dirname "$host")"
            found=1
            
            # 每次扫描后短暂等待
            sleep 1
        fi
    done
    
    if [ $found -eq 0 ]; then
        echo "❌ 没有找到可写的SCSI主机"
        return 1
    fi
    
    echo "✅ 手动SCSI扫描完成"
    return 0
}

# 函数：重置USB控制器
reset_usb_controller() {
    echo "🔌 重置USB控制器..."
    
    # 方法1: 重新加载USB存储驱动
    echo "重新加载USB存储驱动..."
    sudo modprobe -r usb_storage && sleep 2 && sudo modprobe usb_storage
    
    # 方法2: 重置所有USB设备授权
    echo "重置USB设备授权..."
    for usb in /sys/bus/usb/devices/*/authorized; do
        if [ -w "$usb" ]; then
            echo 0 | sudo tee "$usb" >/dev/null 2>&1
            echo 1 | sudo tee "$usb" >/dev/null 2>&1
        fi
    done
    
    # 方法3: 重启USB相关服务
    echo "重启USB相关服务..."
    sudo systemctl restart systemd-udevd 2>/dev/null
}

# 函数：模拟物理重插
simulate_physical_replug() {
    echo "🔧 模拟物理重插过程..."
    
    # 等待设备稳定
    echo "等待系统稳定..."
    sleep 3
    
    # 重新扫描
    rescan_scsi_bus
    
    # 触发udev事件
    echo "触发设备事件..."
    sudo udevadm trigger --subsystem-match=block --action=add
    sudo udevadm trigger --subsystem-match=usb --action=add
    
    # 等待设备识别
    echo "等待设备识别..."
    sleep 5
}

# 函数：检查新设备
check_new_devices() {
    echo "🔍 检查新识别的设备..."
    
    local original_count=$1
    local current_count=$(lsblk -d -o NAME | grep -c '^sd')
    
    if [ $current_count -gt $original_count ]; then
        echo "✅ 发现新设备！"
        echo "📀 当前磁盘列表："
        lsblk -d -o NAME,SIZE,TYPE,MOUNTPOINT | grep -E "^(NAME|sd)"
        return 0
    else
        echo "❌ 未发现新设备"
        return 1
    fi
}

# 主程序开始
echo "🚀 开始设备重新识别流程..."
echo "💡 请确保USB设备已物理插入"

# 记录当前的设备数量
original_device_count=$(lsblk -d -o NAME | grep -c '^sd')
echo "当前磁盘设备数量: $original_device_count"

# 显示USB控制器信息
show_usb_controllers

# 用户确认
read -p "是否继续执行设备重新识别？(y/N): " confirm
if [[ ! $confirm =~ ^[Yy]$ ]]; then
    echo "❌ 用户取消操作"
    exit 0
fi

echo ""
echo "🔄 开始执行重新识别流程..."
echo "========================"

# 执行重新识别步骤
reset_usb_controller
simulate_physical_replug

# 检查结果
echo ""
echo "📊 重新识别结果："
echo "========================"
if check_new_devices $original_device_count; then
    echo ""
    echo "🎉 设备重新识别成功！"
    echo "💡 新设备应该已经在系统中可用"
    
    # 显示详细的设备信息
    echo ""
    echo "📋 详细的设备信息："
    lsblk -o NAME,SIZE,TYPE,MOUNTPOINT,FSTYPE,LABEL | grep -E "^(NAME|sd)"
else
    echo ""
    echo "⚠️  未检测到新设备，可能的原因："
    echo "   - 设备未物理插入"
    echo "   - 设备连接问题"
    echo "   - 设备硬件故障"
    echo "   - 需要手动物理重插"
    
    # 建议手动操作
    echo ""
    echo "💡 建议尝试："
    echo "   1. 物理拔出USB设备"
    echo "   2. 等待5秒"
    echo "   3. 重新插入USB设备"
    echo "   4. 再次运行此脚本"
fi

# 显示系统日志相关信息
echo ""
echo "📝 系统日志相关信息："
sudo dmesg | tail -10 | grep -E "usb|scsi|sd" | tail -5